package com.bst.datax.executor.service.jobhandler;

import com.bst.datax.executor.service.jobhandler.streamload.DataShardingStrategy;
import com.bst.datax.executor.service.jobhandler.streamload.DorisDataConverter;
import com.bst.datax.executor.service.jobhandler.streamload.DorisStreamLoadService;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.jdbc.core.JdbcTemplate;

import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.stream.Collectors;

public class StreamLoadHandler {
    private DorisStreamLoadService streamLoadService;
    private JdbcTemplate dorisJdbcTemplate;
    private DorisDataConverter dataConverter;

    @Qualifier("dorisStreamLoadThreadPool")
    private ExecutorService threadPool;

    // @Retryable(maxAttempts = 3, backoff = @Backoff(delay = 1000))
    public void parallelStreamLoad(String sourceTable, String targetTable, int parallelism) {
        // 1. 获取数据分片条件
        long minId = dorisJdbcTemplate.queryForObject(
                "SELECT MIN(id) FROM " + sourceTable, Long.class);
        long maxId = dorisJdbcTemplate.queryForObject(
                "SELECT MAX(id) FROM " + sourceTable, Long.class);

        List<String> shardConditions = DataShardingStrategy.createIdRangeShards("id", minId, maxId, parallelism);

        // 2. 创建并行任务
        List<CompletableFuture<Void>> futures = shardConditions.stream()
                .map(condition -> CompletableFuture.runAsync(() ->
                        processShard(sourceTable, targetTable, condition), threadPool))
                .collect(Collectors.toList());

        // 3. 等待所有任务完成
        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
    }

    private void processShard(String sourceTable, String targetTable, String condition) {
        try {
            // 1. 查询分片数据
            List<Map<String, Object>> data = dorisJdbcTemplate.queryForList(
                    "SELECT * FROM " + sourceTable + " WHERE " + condition);

            // 2. 转换为CSV格式
            String csvData = dataConverter.toCsv(data);

            // 3. 执行Stream Load
            streamLoadService.streamLoad("your_db", targetTable, csvData, "csv");

            // 4. 记录处理状态
            dorisJdbcTemplate.update(
                    "UPDATE " + sourceTable + " SET is_processed = 1 WHERE " + condition);

        } catch (Exception e) {
            throw new RuntimeException("Process shard failed: " + condition, e);
        }
    }
}
